#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI内容解析模块
支持多种AI服务和解析策略
整合华为云、阿里云、OpenAI等多种AI服务
"""

import os
import sys
import json
import logging
import requests
import re
import time
from typing import Dict, List, Any, Optional, Union
from abc import ABC, abstractmethod

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import LOG_CONFIG

# 配置日志
logging.basicConfig(
    level=getattr(logging, LOG_CONFIG["level"]),
    format=LOG_CONFIG["format"]
)
logger = logging.getLogger(__name__)


class BaseAIProvider(ABC):
    """AI服务提供商基类"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.name = self.__class__.__name__
    
    @abstractmethod
    def call_api(self, prompt: str, **kwargs) -> str:
        """调用AI API"""
        pass
    
    @abstractmethod
    def is_available(self) -> bool:
        """检查服务是否可用"""
        pass


class HuaweiAIProvider(BaseAIProvider):
    """华为云AI服务提供商"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.api_url = config.get('api_url', 'https://your-huawei-api-url')
        self.api_key = config.get('api_key', '')
        self.model = config.get('model', 'default')
    
    def call_api(self, prompt: str, **kwargs) -> str:
        """调用华为云AI API"""
        try:
            headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {self.api_key}'
            }
            
            payload = {
                'model': self.model,
                'messages': [
                    {'role': 'user', 'content': prompt}
                ],
                'temperature': kwargs.get('temperature', 0.7),
                'max_tokens': kwargs.get('max_tokens', 2000)
            }
            
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            
            data = response.json()
            return data.get('choices', [{}])[0].get('message', {}).get('content', '')
            
        except Exception as e:
            logger.error(f"华为云AI API调用失败: {e}")
            raise
    
    def is_available(self) -> bool:
        """检查华为云服务是否可用"""
        try:
            test_prompt = "测试连接"
            response = self.call_api(test_prompt, max_tokens=10)
            return bool(response)
        except:
            return False


class OpenAIProvider(BaseAIProvider):
    """OpenAI服务提供商"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.api_url = config.get('api_url', 'https://api.openai.com/v1/chat/completions')
        self.api_key = config.get('api_key', '')
        self.model = config.get('model', 'gpt-3.5-turbo')
    
    def call_api(self, prompt: str, **kwargs) -> str:
        """调用OpenAI API"""
        try:
            headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {self.api_key}'
            }
            
            payload = {
                'model': self.model,
                'messages': [
                    {'role': 'user', 'content': prompt}
                ],
                'temperature': kwargs.get('temperature', 0.7),
                'max_tokens': kwargs.get('max_tokens', 2000)
            }
            
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            
            data = response.json()
            return data.get('choices', [{}])[0].get('message', {}).get('content', '')
            
        except Exception as e:
            logger.error(f"OpenAI API调用失败: {e}")
            raise
    
    def is_available(self) -> bool:
        """检查OpenAI服务是否可用"""
        try:
            test_prompt = "Hello"
            response = self.call_api(test_prompt, max_tokens=10)
            return bool(response)
        except:
            return False


class DifyAIProvider(BaseAIProvider):
    """Dify AI服务提供商"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.api_url = config.get('api_url', '')
        self.api_key = config.get('api_key', '')
    
    def call_api(self, prompt: str, **kwargs) -> str:
        """调用Dify AI API"""
        try:
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            payload = {
                'inputs': {'query': prompt},
                'response_mode': 'blocking',
                'user': kwargs.get('user', 'spider-user')
            }
            
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            
            data = response.json()
            return data.get('answer', '')
            
        except Exception as e:
            logger.error(f"Dify AI API调用失败: {e}")
            raise
    
    def is_available(self) -> bool:
        """检查Dify服务是否可用"""
        try:
            test_prompt = "测试"
            response = self.call_api(test_prompt)
            return bool(response)
        except:
            return False


class MockAIProvider(BaseAIProvider):
    """模拟AI服务提供商（用于测试）"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
    
    def call_api(self, prompt: str, **kwargs) -> str:
        """模拟AI API调用"""
        logger.info("使用模拟AI服务进行内容解析...")
        
        # 简单的模拟响应
        if "标题" in prompt:
            return '{"标题": "模拟标题", "正文": "这是模拟的正文内容", "作者": "模拟作者", "发布时间": "2024-01-01"}'
        else:
            return "这是一个模拟的AI响应"
    
    def is_available(self) -> bool:
        """模拟服务总是可用"""
        return True


class AIProviderManager:
    """AI服务提供商管理器"""
    
    def __init__(self):
        self.providers: Dict[str, BaseAIProvider] = {}
        self.fallback_order = ['huawei', 'openai', 'dify', 'mock']
        self.current_provider = None
    
    def register_provider(self, name: str, provider: BaseAIProvider):
        """注册AI服务提供商"""
        self.providers[name] = provider
        logger.info(f"注册AI服务提供商: {name}")
    
    def get_available_provider(self) -> Optional[BaseAIProvider]:
        """获取可用的AI服务提供商"""
        if self.current_provider and self.current_provider.is_available():
            return self.current_provider
        
        # 按优先级查找可用的服务
        for provider_name in self.fallback_order:
            if provider_name in self.providers:
                provider = self.providers[provider_name]
                if provider.is_available():
                    self.current_provider = provider
                    logger.info(f"切换到AI服务提供商: {provider_name}")
                    return provider
        
        logger.error("没有可用的AI服务提供商")
        return None
    
    def call_ai(self, prompt: str, **kwargs) -> str:
        """调用AI服务"""
        provider = self.get_available_provider()
        if not provider:
            raise Exception("没有可用的AI服务提供商")
        
        try:
            return provider.call_api(prompt, **kwargs)
        except Exception as e:
            logger.error(f"AI服务调用失败: {e}")
            # 尝试切换到下一个提供商
            self.current_provider = None
            provider = self.get_available_provider()
            if provider and provider != self.current_provider:
                return provider.call_api(prompt, **kwargs)
            raise


class ContentFieldExtractor:
    """内容字段提取器"""
    
    def __init__(self):
        self.field_templates = {
            'news': {
                'fields': ['标题', '正文', '作者', '发布时间', '来源', '摘要'],
                'system_prompt': '你是一个专业的新闻内容提取专家，请准确提取新闻文章的各个字段。'
            },
            'article': {
                'fields': ['标题', '正文', '作者', '发布时间', '分类', '标签'],
                'system_prompt': '你是一个专业的文章内容提取专家，请准确提取文章的各个字段。'
            },
            'product': {
                'fields': ['产品名称', '价格', '描述', '品牌', '规格', '评分'],
                'system_prompt': '你是一个专业的商品信息提取专家，请准确提取商品的各个字段。'
            },
            'financial': {
                'fields': ['公司名称', '营业收入', '净利润', '发布日期', '报告期间'],
                'system_prompt': '你是一个专业的财务信息提取专家，请准确提取财务报告的各个字段。'
            }
        }
    
    def get_template(self, content_type: str) -> Dict[str, Any]:
        """获取内容模板"""
        return self.field_templates.get(content_type, self.field_templates['news'])


class AIContentParser:
    """AI内容解析器主类"""
    
    def __init__(self, config: Optional[Dict] = None):
        self.config = config or {}
        self.provider_manager = AIProviderManager()
        self.field_extractor = ContentFieldExtractor()
        
        # 初始化AI服务提供商
        self._init_providers()
        
        logger.info("AI内容解析器初始化完成")
    
    def _init_providers(self):
        """初始化AI服务提供商"""
        providers_config = self.config.get('providers', {})
        
        # 华为云AI
        if 'huawei' in providers_config:
            huawei_provider = HuaweiAIProvider(providers_config['huawei'])
            self.provider_manager.register_provider('huawei', huawei_provider)
        
        # OpenAI
        if 'openai' in providers_config:
            openai_provider = OpenAIProvider(providers_config['openai'])
            self.provider_manager.register_provider('openai', openai_provider)
        
        # Dify AI
        if 'dify' in providers_config:
            dify_provider = DifyAIProvider(providers_config['dify'])
            self.provider_manager.register_provider('dify', dify_provider)
        
        # 模拟AI（总是注册，作为最后的fallback）
        mock_provider = MockAIProvider({})
        self.provider_manager.register_provider('mock', mock_provider)
    
    def parse_content(self, text: str, content_type: str = 'news', 
                     fields: Optional[List[str]] = None,
                     system_prompt: Optional[str] = None,
                     **kwargs) -> Dict[str, Any]:
        """
        解析内容
        
        Args:
            text: 需要解析的文本
            content_type: 内容类型 (news, article, product, financial)
            fields: 需要提取的字段列表
            system_prompt: 自定义系统提示词
            **kwargs: 其他参数
            
        Returns:
            解析结果字典
        """
        try:
            # 获取模板
            template = self.field_extractor.get_template(content_type)
            
            # 使用自定义字段或模板字段
            target_fields = fields or template['fields']
            prompt_system = system_prompt or template['system_prompt']
            
            # 构建提示词
            prompt = self._build_prompt(text, target_fields, prompt_system)
            
            # 调用AI服务
            response = self.provider_manager.call_ai(prompt, **kwargs)
            
            # 解析响应
            result = self._parse_response(response, target_fields)
            
            return result
            
        except Exception as e:
            logger.error(f"AI内容解析失败: {e}")
            return {field: None for field in (fields or ['标题', '正文'])}
    
    def _build_prompt(self, text: str, fields: List[str], system_prompt: str) -> str:
        """构建AI提示词"""
        field_desc = '、'.join(fields)
        
        # 限制文本长度
        max_text_length = self.config.get('max_text_length', 3000)
        if len(text) > max_text_length:
            text = text[:max_text_length] + "..."
        
        prompt = f"""{system_prompt}

请从以下文本中提取 {field_desc} 字段，以JSON格式返回。

文本内容:
{text}

请按照以下格式返回JSON:
{{
    "{fields[0]}": "提取的内容1",
    "{fields[1]}": "提取的内容2"
}}

提取要求:
1. 如果某个字段在文本中不存在或无法确定，请返回null
2. 正文内容需要提取完整，但要去除页眉页尾等无关信息
3. 时间格式请统一为YYYY-MM-DD HH:MM:SS格式
4. 数字金额请保留原始格式
5. 只返回JSON格式，不要添加其他解释文字
"""
        return prompt
    
    def _parse_response(self, response: str, fields: List[str]) -> Dict[str, Any]:
        """解析AI响应"""
        try:
            # 清理响应文本
            response = response.strip()
            
            # 尝试提取JSON
            json_match = re.search(r'\{.*?\}', response, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
                result = json.loads(json_str)
            else:
                # 如果没有找到JSON格式，尝试解析为简单文本
                logger.warning("AI响应不是JSON格式，尝试文本解析")
                result = self._parse_text_response(response, fields)
            
            # 验证和清理结果
            parsed_result = {}
            for field in fields:
                value = result.get(field)
                if value and isinstance(value, str):
                    value = value.strip()
                    if value.lower() in ['null', 'none', '无', '']:
                        value = None
                parsed_result[field] = value
            
            return parsed_result
            
        except Exception as e:
            logger.error(f"解析AI响应失败: {e}")
            logger.debug(f"原始响应: {response}")
            return {field: None for field in fields}
    
    def _parse_text_response(self, response: str, fields: List[str]) -> Dict[str, Any]:
        """解析文本格式的响应"""
        result = {}
        
        for field in fields:
            # 简单的文本匹配
            pattern = rf'{field}[：:]\s*(.+?)(?=\n|$)'
            match = re.search(pattern, response)
            if match:
                result[field] = match.group(1).strip()
            else:
                result[field] = None
        
        return result
    
    def batch_parse(self, texts: List[str], content_type: str = 'news',
                   **kwargs) -> List[Dict[str, Any]]:
        """批量解析内容"""
        results = []
        
        for i, text in enumerate(texts):
            logger.info(f"批量解析进度: {i+1}/{len(texts)}")
            
            try:
                result = self.parse_content(text, content_type, **kwargs)
                results.append(result)
                
                # 添加延迟避免API限制
                time.sleep(self.config.get('batch_delay', 1))
                
            except Exception as e:
                logger.error(f"批量解析第{i+1}个文本失败: {e}")
                results.append({})
        
        return results
    
    def get_provider_status(self) -> Dict[str, bool]:
        """获取所有AI服务提供商的状态"""
        status = {}
        for name, provider in self.provider_manager.providers.items():
            status[name] = provider.is_available()
        return status


# 使用示例
def main():
    """主函数"""
    # 配置AI服务
    config = {
        'providers': {
            'huawei': {
                'api_url': 'https://your-huawei-api-url',
                'api_key': 'your-api-key',
                'model': 'your-model'
            },
            'openai': {
                'api_url': 'https://api.openai.com/v1/chat/completions',
                'api_key': 'your-openai-key',
                'model': 'gpt-3.5-turbo'
            }
        },
        'max_text_length': 3000,
        'batch_delay': 1
    }
    
    # 创建解析器
    parser = AIContentParser(config)
    
    # 测试文本
    test_text = """
    标题：人工智能技术发展现状与趋势
    
    近年来，人工智能技术发展迅速，在各个领域都有广泛应用。
    本文分析了当前AI技术的发展现状，并对未来趋势进行了预测。
    
    作者：张三
    发布时间：2024-01-15 10:30:00
    来源：科技日报
    """
    
    # 解析内容
    result = parser.parse_content(
        text=test_text,
        content_type='news',
        fields=['标题', '正文', '作者', '发布时间', '来源']
    )
    
    print("解析结果:")
    for key, value in result.items():
        print(f"{key}: {value}")
    
    # 检查服务状态
    status = parser.get_provider_status()
    print("\nAI服务状态:")
    for provider, available in status.items():
        print(f"{provider}: {'可用' if available else '不可用'}")


if __name__ == "__main__":
    main() 